# Copyright (c) 2021 Presto Labs Pte. Ltd.
# Author: hjkim
# pylint: disable=wrong-import-position,too-few-public-methods
# pylint: disable=unnecessary-comprehension,too-many-locals,broad-except
# pylint: disable=no-name-in-module
import os
from os.path import join as pjoin
import time
import json
import datetime
import tempfile
import pandas
import copy
from collections import defaultdict

import fire

import python.coin_deploy.lm_agg2.kline_util as kline_util
from google.protobuf.json_format import MessageToDict
from cc.appcoin2.strategy.lm.kline_dumper_live import get_kline_request_dict_and_vu_rank, get_symbol_universe, is_enable_clear_all_1, intv_to_seconds

from xunkemgmt_client.tool.slack_noti import send_to_slack

os.environ["TZ"] = "GMT"
time.tzset()

def run(now=None, kline_path=None, test=False, exchange="Binance",
        use_slack=False, tolerance=100, config_path=None, lm_strat_json_from_conf_gen=None):
  assert config_path is not None
  if is_enable_clear_all_1(config_path):
    with open(kline_path, "w", encoding="utf8") as fout:
      fout.write(json.dumps({"klines": []}, indent=2))
    return

  import cc.appcoin2.strategy.lm.kline_generator as klg
  import cc.appcoin2.strategy.lm.kline_dumper_live_aux as kline_aux
  kline_request_dict, vu_rank = get_kline_request_dict_and_vu_rank(config_path)

  el_st = time.time()
  if now is None:
    now = datetime.datetime.now()
  elif isinstance(now, (int, str)):
    now = datetime.datetime.strptime(str(now), "%Y%m%d%H%M")
  else:
    raise ValueError(now)

  sequential = False
  if "sequential" in kline_request_dict:
    sequential = bool(kline_request_dict.pop("sequential"))
  dump_ref = False
  if "dump_ref" in kline_request_dict:
    dump_ref = bool(kline_request_dict.pop("dump_ref"))
  print("Dump Ref: ", dump_ref)

  kline_request_dict = {
      key: kline_request_dict[key]
      for key in sorted(kline_request_dict.keys(), key=lambda x: intv_to_seconds[x])}

  print("kline period info:")
  start_dt_dict = {}
  end_dt_dict = {}
  end_dt_it = copy.deepcopy(now)
  for intv, num in kline_request_dict.items():
    start_dt_dict[intv] = end_dt_it - datetime.timedelta(seconds=intv_to_seconds[intv]*num)
    end_dt_dict[intv] = end_dt_it
    if sequential:
      end_dt_it = start_dt_dict[intv]

  kline_request_dict = {
      key: kline_request_dict[key]
      for key in sorted(kline_request_dict.keys(), key=lambda x: intv_to_seconds[x], reverse=True)}
  for intv, num in kline_request_dict.items():
    print(f"\tstart:\t{start_dt_dict[intv]}")
    print(f"\tend:\t{end_dt_dict[intv]} ({intv} * {num})")

  if dump_ref:
    print("dump_ref: on")

  volume_info = None
  if lm_strat_json_from_conf_gen is not None and vu_rank:
    if vu_rank == "all":
      vu_rank_ge, vu_rank_lt = 0, 10000
    else:
      vu_rank_ge, vu_rank_lt = list(map(int, vu_rank.split("-"))) if "-" in vu_rank else [0, int(vu_rank)]
    volume_info = json.loads(open(lm_strat_json_from_conf_gen).read()).get("volume_info", {})

  symbol_map, focus_exchange, ref_exchanges = get_symbol_universe(config_path, volume_info)
  focus_market_type = "Futures" if "Futures" in focus_exchange else "Spot"
  ref_market_types = ["Futures" if "Futures" in ref_exchange else "Spot" for ref_exchange in ref_exchanges]
  focus_exchange = focus_exchange.replace("Futures", "")
  ref_exchanges = [ref_exchange.replace("Futures", "") for ref_exchange in ref_exchanges]
  focus_mea = f"{focus_market_type}.{focus_exchange}.{kline_util.api_ver_map[focus_exchange][focus_market_type]}"
  ref_meas = [f"{ref_market_type}.{ref_exchange}.{kline_util.api_ver_map[ref_exchange][ref_market_type]}"
      for ref_market_type, ref_exchange in zip(ref_market_types, ref_exchanges)]

  # filter symbols by vu logic
  if volume_info is not None:
    volume_info_by_exchs = volume_info.get("volume_info_by_exchs", [])
    symbol_map_whole = copy.deepcopy(symbol_map)
    symbol_map = {}
    if len(volume_info_by_exchs) > 0:
      volume_info_list = sorted([
          (volume_info_entry["volume"], volume_info_entry["symbol"])
          for volume_info_entry in volume_info_by_exchs[0].get("volume_infos", [])],
          reverse=True)
      cnt = 0
      print(pandas.DataFrame([{"focus": key, "ref": value} for key, value in symbol_map_whole.items()]).to_string())
      print(pandas.DataFrame(volume_info_list).to_string())
      for _, symbol in volume_info_list:
        if symbol_map_whole[symbol]:
          if dump_ref and (None in symbol_map_whole[symbol]):
            continue
          if vu_rank_ge <= cnt < vu_rank_lt or "BTC-" in symbol or "ETH-" in symbol:
            symbol_map[symbol] = symbol_map_whole[symbol]
          cnt += 1

  focus_product0 = list(symbol_map.keys())[0]
  ref_product0s = symbol_map[focus_product0]
  focus_mea = f"{focus_mea}{kline_util.get_api_suffix(focus_market_type, focus_exchange, focus_product0)}"
  ref_meas = [f"{ref_mea}{kline_util.get_api_suffix(ref_mea.split('.')[0], ref_mea.split('.')[1], ref_product0s[i])}" for i, ref_mea in enumerate(ref_meas)]

  kline_info = {"klines": []}
  failed_info = {"failed": []}

  def convert_klinelist(klinelist):
    return MessageToDict(klinelist,
                        including_default_value_fields=True,
                        preserving_proto_field_name=True)

  focus_symbols = []
  ref_symbols_list = [[] for _ in range(len(ref_exchanges))]
  for focus_symbol, ref_symbols in symbol_map.items():
    focus_symbols.append(focus_symbol)
    for i, ref_symbol in enumerate(ref_symbols):
      ref_symbols_list[i].append(ref_symbol)

  for intv, num in kline_request_dict.items():
    adjusted_start_dt = {}

    focus_klines_from_interval_feed = defaultdict(list)
    ref_klines_list_from_interval_feed = [defaultdict(list) for _ in range(len(ref_exchanges))]
    if intv == "1m":
      focus_klines = kline_aux.get_kline_from_interval_feed(start_dt_dict[intv], end_dt_dict[intv], focus_mea, focus_symbols)
      ref_klines_list = [kline_aux.get_kline_from_interval_feed(start_dt_dict[intv], end_dt_dict[intv], ref_mea, ref_symbols)
          for ref_mea, ref_symbols in zip(ref_meas, ref_symbols_list)]
      
      for symb in focus_klines.keys():
        if len(focus_klines[symb]) > 0:
          first_data_dt = pandas.DatetimeIndex([focus_klines[symb][0]["kline_timestamp"]]).to_pydatetime()[0]
          if first_data_dt <= start_dt_dict[intv] + datetime.timedelta(minutes=10):
            last_data_dt = pandas.DatetimeIndex([focus_klines[symb][-1]["kline_timestamp"]]).to_pydatetime()[0]
            adjusted_start_dt[f"{focus_market_type}:{focus_exchange}:{symb}"] = last_data_dt + datetime.timedelta(minutes=1)
          else:
            focus_klines[symb] = []
      focus_klines_from_interval_feed = focus_klines

      for ref_market_type, ref_exchange, ref_klines in zip(ref_market_types, ref_exchanges, ref_klines_list):
        for symb in ref_klines.keys():
          if len(ref_klines[symb]) > 0:
            first_data_dt = pandas.DatetimeIndex([ref_klines[symb][0]["kline_timestamp"]]).to_pydatetime()[0]
            if first_data_dt <= start_dt_dict[intv] + datetime.timedelta(minutes=10):
              last_data_dt = pandas.DatetimeIndex([ref_klines[symb][-1]["kline_timestamp"]]).to_pydatetime()[0]
              adjusted_start_dt[f"{ref_market_type}:{ref_exchange}:{symb}"] = last_data_dt + datetime.timedelta(minutes=1)
            else:
              ref_klines[symb] = []
      ref_klines_list_from_interval_feed = ref_klines_list

    for focus_symbol, ref_symbols in symbol_map.items():
      mep = f"{focus_market_type}:{focus_exchange}:{focus_symbol}"
      # for debug
      if intv == "1m" and mep not in adjusted_start_dt:
        print(f"{mep} will be fully queried from rest api")
      # for debug
      focus_klinelist, focus_failed = klg.query_kline_as_pb(
          now,
          mep,
          adjusted_start_dt[mep] if mep in adjusted_start_dt else start_dt_dict[intv],
          end_dt_dict[intv], kline_period=intv, max_cnt=1000, tolerance=tolerance)
      focus_klinelist = convert_klinelist(focus_klinelist)
      failed_info["failed"] += focus_failed
      ohlc = focus_klines_from_interval_feed[focus_symbol] + focus_klinelist["klines"]
      kline_info["klines"].append({
        "symbol": f"{focus_exchange}{focus_market_type.replace('Spot', '')}_{focus_symbol}",
        "interval_ns": int(intv_to_seconds[intv] * 1e9),
        "ohlc": ohlc})
      print(f"kline of {focus_symbol}(Focus) ({intv}) fetched")
      if dump_ref:
        for ref_symbol, ref_market_type, ref_exchange, ref_klines_from_interval_feed in zip(ref_symbols, ref_market_types, ref_exchanges, ref_klines_list_from_interval_feed):
          if ref_symbol is None:
            continue
          mep = f"{ref_market_type}:{ref_exchange}:{ref_symbol}"
          # for debug
          if intv == "1m" and mep not in adjusted_start_dt:
            print(f"{mep} will be fully queried from rest api")
          # for debug
          ref_klinelist, ref_failed = klg.query_kline_as_pb(
              now,
              mep,
              adjusted_start_dt[mep] if mep in adjusted_start_dt else start_dt_dict[intv],
              end_dt_dict[intv], kline_period=intv, tolerance=tolerance)
          ref_klinelist = convert_klinelist(ref_klinelist)
          failed_info["failed"] += ref_failed
          ohlc = ref_klines_from_interval_feed[ref_symbol] + ref_klinelist["klines"]
          kline_info["klines"].append({
            "symbol": f"{ref_exchange}{ref_market_type.replace('Spot', '')}_{ref_symbol}",
            "interval_ns": int(intv_to_seconds[intv] * 1e9),
            "ohlc": ohlc
          })
          print(f"kline of {ref_symbol}(Ref) ({intv}) fetched")
  dt_str = now.strftime("%Y%m%d%H")
  kline_path = kline_path or \
    pjoin(tempfile.mkdtemp(prefix=f"coin_midfreq_kline_{dt_str}", dir="/tmp"),
          "kline_info.json")
  print(f"dump kline to {kline_path}")
  with open(kline_path, "w", encoding="utf8") as fout:
    fout.write(json.dumps(kline_info, indent=2))
  elapsed = time.time() - el_st
  print(f"elapsed to get kline: {elapsed:.4f} sec")
  if use_slack:
    channel = "#k_hft_operation_log"
    if failed_info["failed"]:
      failed_df = pandas.DataFrame(failed_info["failed"])
      failed_cnt = failed_df['product'].value_counts().sort_values()
      send_to_slack('\n'.join([
                      f'Hostname: {os.environ.get("HOSTNAME", "")}',
                      f'kline_path: {kline_path}',
                      "",
                      "Failure counts",
                      failed_cnt.to_string(),
                      "",
                      "Failures",
                      failed_df.to_string()]),
                      channel,
                      'file',
                      now.strftime(f"Kline failures %Y%m%d_%H%M%S {focus_exchange.replace('Futures', '')}"))


if __name__ == "__main__":
  fire.Fire()
